package com.xiaozhou.controller;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.xiaozhou.function.FunctionManager;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author shhzhou
 * @description: Agent模型 例子
 * @date 2024/11/1 9:07
 */
@RestController
@RequestMapping("/agent")
public class AgentController {

    @Autowired
    private ClientV4 clientV4;

    @Autowired
    private FunctionManager functionManager;

    /**
     * @return
     * @return reactor.core.publisher.Flux<java.lang.String>
     * @Author xiaozhou
     * @Description type: function 范例    注意：函数形式调用始终是需要手动调用函数，并将函数结果以聊天形式重新发送给大模型进行内容整合
     * @Date 9:14 2024/11/4
     * @Param
     */
    @RequestMapping("/functionModel")
    public Flux<String> functionModel() {

        // 构造消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();
        Map<String, String> messageContent = new HashMap<>();
        messageContent.put("type", "text");
        messageContent.put("text", "帮我查询2018年至2023年，每年五一假期全国旅游出行数据，并绘制成柱状图展示数据趋势。");
        chatMessages.add(new ChatMessage("user", Collections.singletonList(messageContent)));

        //构造函数列表
        List<ChatTool> toolList = new ArrayList<>();

        ChatTool tool = new ChatTool();
        tool.setType(ChatToolType.FUNCTION.value());

        ChatFunction chatFunction = new ChatFunction();
        chatFunction.setName("getTravelData");
        chatFunction.setDescription("获取每年的全国旅游数据");

        ChatFunctionParameters chatFunctionParameters = new ChatFunctionParameters();
        chatFunctionParameters.setType("object");

        Map<String, Object> properties = new HashMap<>();
        properties.put("fromYear", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "开始年份，格式为yyyy");
        }});
        properties.put("toYear", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "结束年份，格式为yyyy");
        }});
        chatFunctionParameters.setProperties(properties);
        chatFunction.setParameters(chatFunctionParameters);
        chatFunction.setRequired(Arrays.asList("fromYear", "toYear"));
        tool.setFunction(chatFunction);
        toolList.add(tool);

        // 先获取大模型要调用的函数信息
        ChatCompletionRequest toolRequest = new ChatCompletionRequest();
        toolRequest.setModel("glm-4-alltools");
        toolRequest.setMessages(chatMessages);
        toolRequest.setStream(true);
        toolRequest.setTools(toolList);
        toolRequest.setInvokeMethod(Constants.invokeMethod);

        ModelApiResponse toolModelApiResponse = clientV4.invokeModelApi(toolRequest);


        Flowable<ModelData> flowable = toolModelApiResponse.getFlowable();

        AtomicReference<String> result = new AtomicReference<>();

        flowable.blockingSubscribe(modelData -> {
            if (modelData.getChoices().get(0).getFinishReason() != null && modelData.getChoices().get(0).getFinishReason().equals("tool_calls")) {
                String name = modelData.getChoices().get(0).getDelta().getTool_calls().get(0).getFunction().getName();
                String fromYear = new ObjectMapper().readTree(modelData.getChoices().get(0).getDelta().getTool_calls().get(0).getFunction().getArguments().textValue()).findValue("fromYear").asText();
                String toYear = new ObjectMapper().readTree(modelData.getChoices().get(0).getDelta().getTool_calls().get(0).getFunction().getArguments().textValue()).findValue("toYear").asText();
                List<String> list = Arrays.asList(fromYear, toYear);
                // 调用函数，获取函数结果
                result.set(functionManager.callFunction(name, list));
            }
        });

        chatMessages.add(new ChatMessage("tool", result));


        // 构造模型对象
        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel("glm-4-alltools");
        request.setMessages(chatMessages);
        request.setStream(true);
        request.setInvokeMethod(Constants.invokeMethod);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(request);


        // 使用 Flux.create 创建流
        return Flux.create(sink -> {
            modelApiResponse.getFlowable().map(modelData -> {
                return new ChatMessageAccumulator(
                        modelData.getChoices().get(0).getDelta(),
                        null,
                        modelData.getChoices().get(0),
                        modelData.getUsage(),
                        modelData.getCreated(),
                        modelData.getId()
                );
            }).doOnNext(chatMessageAccumulator -> {
                String content = chatMessageAccumulator.getDelta().getContent();
                if (content != null) {
                    sink.next(content);
                }
            }).doOnComplete(() -> {
                sink.complete();
            }).blockingSubscribe();
        });

    }


    /**
     * @return
     * @return reactor.core.publisher.Flux<java.lang.String>
     * @Author xiaozhou
     * @Description type: code_interpreter 范例
     * @Date 9:33 2024/11/4
     * @Param
     */
    @RequestMapping("/codeInterpreterModel")
    public Flux<ChatMessageAccumulator> codeInterpreterModel() {

        // 构造消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();
        Map<String, String> messageContent = new HashMap<>();
        messageContent.put("type", "text");
        messageContent.put("text", "帮我查询2018年至2023年，每年五一假期全国旅游出行数据，并绘制成柱状图展示数据趋势。");
        chatMessages.add(new ChatMessage("user", Collections.singletonList(messageContent)));

        //构造函数列表
        List<ChatTool> toolList = new ArrayList<>();
        ChatTool tool = new ChatTool();
        tool.setType("code_interpreter");
        toolList.add(tool);

        // 先获取大模型要调用的函数信息
        ChatCompletionRequest toolRequest = new ChatCompletionRequest();
        toolRequest.setModel("glm-4-alltools");
        toolRequest.setMessages(chatMessages);
        toolRequest.setStream(true);
        toolRequest.setTools(toolList);
        toolRequest.setInvokeMethod(Constants.invokeMethod);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(toolRequest);

        // 使用 Flux.create 创建流
        return Flux.create(sink -> {
            modelApiResponse.getFlowable().map(modelData -> {
                return new ChatMessageAccumulator(
                        modelData.getChoices().get(0).getDelta(),
                        null,
                        modelData.getChoices().get(0),
                        modelData.getUsage(),
                        modelData.getCreated(),
                        modelData.getId()
                );
            }).doOnNext(chatMessageAccumulator -> {
                sink.next(chatMessageAccumulator);
            }).doOnComplete(() -> {
                sink.complete();
            }).blockingSubscribe();
        });

    }


    /**
     * @return
     * @return reactor.core.publisher.Flux<com.zhipu.oapi.service.v4.model.ChatMessageAccumulator>
     * @Author xiaozhou
     * @Description type: drawing_tool 范例    但从返回值中无法看到调用类型为drawing_tool，暂不清楚原因
     * @Date 9:47 2024/11/4
     * @Param
     */
    @RequestMapping("/drawingToolModel")
    public Flux<ChatMessageAccumulator> drawingToolModel() {

        // 构造消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();
        Map<String, String> messageContent = new HashMap<>();
        messageContent.put("type", "text");
        messageContent.put("text", "帮我查询2018年至2023年，每年五一假期全国旅游出行数据，并绘制成柱状图展示数据趋势。");
        chatMessages.add(new ChatMessage("user", Collections.singletonList(messageContent)));

        //构造函数列表
        List<ChatTool> toolList = new ArrayList<>();
        ChatTool tool = new ChatTool();
        tool.setType("drawing_tool");
        toolList.add(tool);

        // 先获取大模型要调用的函数信息
        ChatCompletionRequest toolRequest = new ChatCompletionRequest();
        toolRequest.setModel("glm-4-alltools");
        toolRequest.setMessages(chatMessages);
        toolRequest.setStream(true);
        toolRequest.setTools(toolList);
        toolRequest.setInvokeMethod(Constants.invokeMethod);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(toolRequest);

        // 使用 Flux.create 创建流
        return Flux.create(sink -> {
            modelApiResponse.getFlowable().map(modelData -> {
                return new ChatMessageAccumulator(
                        modelData.getChoices().get(0).getDelta(),
                        null,
                        modelData.getChoices().get(0),
                        modelData.getUsage(),
                        modelData.getCreated(),
                        modelData.getId()
                );
            }).doOnNext(chatMessageAccumulator -> {
                sink.next(chatMessageAccumulator);
            }).doOnComplete(() -> {
                sink.complete();
            }).blockingSubscribe();
        });

    }


    /**
     * @return
     * @return reactor.core.publisher.Flux<com.zhipu.oapi.service.v4.model.ChatMessageAccumulator>
     * @Author xiaozhou
     * @Description type: web_browser 范例
     * @Date 9:53 2024/11/4
     * @Param
     */
    @RequestMapping("/webBrowserModel")
    public Flux<ChatMessageAccumulator> webBrowserModel() {

        // 构造消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();
        Map<String, String> messageContent = new HashMap<>();
        messageContent.put("type", "text");
        messageContent.put("text", "帮我查询2018年至2023年，每年五一假期全国旅游出行数据，并绘制成柱状图展示数据趋势。");
        chatMessages.add(new ChatMessage("user", Collections.singletonList(messageContent)));

        //构造函数列表
        List<ChatTool> toolList = new ArrayList<>();
        ChatTool tool = new ChatTool();
        tool.setType("web_browser");
        toolList.add(tool);

        // 先获取大模型要调用的函数信息
        ChatCompletionRequest toolRequest = new ChatCompletionRequest();
        toolRequest.setModel("glm-4-alltools");
        toolRequest.setMessages(chatMessages);
        toolRequest.setStream(true);
        toolRequest.setTools(toolList);
        toolRequest.setInvokeMethod(Constants.invokeMethod);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(toolRequest);

        // 使用 Flux.create 创建流
        return Flux.create(sink -> {
            modelApiResponse.getFlowable().map(modelData -> {
                return new ChatMessageAccumulator(
                        modelData.getChoices().get(0).getDelta(),
                        null,
                        modelData.getChoices().get(0),
                        modelData.getUsage(),
                        modelData.getCreated(),
                        modelData.getId()
                );
            }).doOnNext(chatMessageAccumulator -> {
                sink.next(chatMessageAccumulator);
            }).doOnComplete(() -> {
                sink.complete();
            }).blockingSubscribe();
        });

    }


    /**
     * @return
     * @return reactor.core.publisher.Flux<com.zhipu.oapi.service.v4.model.Delta>
     * @Author xiaozhou
     * @Description 流式调用模型   该方法更好，可以灵活调控返回值，只需调节返回值类型和next方法中的参数即可
     * @Date 14:56 2024/11/1
     * @Param
     */
    @RequestMapping("/streamModel")
    public Flux<Delta> streamModel(String message) {

        //构造消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();
        // 设置默认系统文本（简单理解是构建模型的人设）
        chatMessages.add(new ChatMessage("system", "你的身份是一个专业的市场营销专家，你的名字叫小林，请以恭敬的语气回答客户的问题！"));
        chatMessages.add(new ChatMessage("user", message));

        // 构造模型对象
        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel("glm-4-flash");
        request.setMessages(chatMessages);
        request.setStream(true);
        request.setInvokeMethod(Constants.invokeMethod);

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(request);

        // 使用 Flux.create 创建流
        return Flux.create(sink -> {
            modelApiResponse.getFlowable().map(modelData -> {
                return new ChatMessageAccumulator(
                        modelData.getChoices().get(0).getDelta(),
                        null,
                        modelData.getChoices().get(0),
                        modelData.getUsage(),
                        modelData.getCreated(),
                        modelData.getId()
                );
            }).doOnNext(chatMessageAccumulator -> {
                sink.next(chatMessageAccumulator.getDelta());
            }).doOnComplete(() -> {
                sink.complete();
            }).blockingSubscribe();
        });

    }


}
